import pandas as pd  # 导入pandas库用于数据处理
from datetime import datetime, timedelta  # 导入datetime和timedelta用于日期计算

def calculate_top_n_changes(csv_file_path, M=5, N=10, ascending=False, X=None):
    # 读取CSV文件
    df = pd.read_csv(csv_file_path)  # 使用pd.read_csv函数读取CSV文件

    # 将trade_date转换为日期格式
    df['trade_date'] = pd.to_datetime(df['trade_date'], format='%Y%m%d')  # 将trade_date列转换为日期格式

    # 如果指定了X年，则过滤出最近X年的数据
    if X is not None:
        current_date = datetime.now()
        start_date = current_date - timedelta(days=X * 365)  # 近似计算X年前的日期
        df = df[df['trade_date'] >= start_date]  # 过滤出最近X年的数据

    # 按ts_code和trade_date排序
    df.sort_values(by=['ts_code', 'trade_date'], inplace=True)  # 按ts_code和trade_date排序

    # 计算过去M天的收盘价变化率
    df['close_change_pct'] = df.groupby('ts_code')['close'].pct_change(M)  # 使用groupby和pct_change函数计算每个股票在过去M天内的收盘价变化率
    # 变化率绝对值
    df['close_change_pct_abs'] = df['close_change_pct'].abs()

    # 找出每个股票在过去M天内的最大涨幅或跌幅的索引
    def get_top_indices(series, N, ascending):
        if ascending:
            return series.nsmallest(N).index
        else:
            return series.nlargest(N).index

    top_indices = df.groupby('ts_code')['close_change_pct_abs'].apply(get_top_indices, N=N, ascending=ascending)

    # 展平多级索引
    top_indices = top_indices.explode()

    # 根据索引提取结果
    result_df = df.loc[top_indices].reset_index(drop=True)

    return result_df  # 返回结果DataFrame


# 示例调用
if __name__ == "__main__":
    M = 25  # 连续几天
    N = 20  # Top N
    ascending = False  # False表示找涨幅，True表示找跌幅
    X = 10 # 最近10年的数据

    # 指数列表
    # 000016.SH,上证50,SSE,中证指数有限公司,规模指数,20031231,1000.0,20040102
    # 000300.SH,沪深300,SSE,中证指数有限公司,规模指数,20041231,1000.0,20050408
    # 000905.SH,中证500,SSE,中证指数有限公司,规模指数,20041231,1000.0,20070115
    # 000688.SH,科创50,SSE,中证指数有限公司,规模指数,20191231,1000.0,20200723
    # 399006.SZ,创业板指,SZSE,深圳证券交易所,规模指数,20100531,1000.0,20100601
    # 399330.SZ,深证100,SZSE,深圳证券交易所,规模指数,20021231,1000.0,20060124
    print(f"最近{X}年，连续{M}天，top{N}的变化率")

    aggregated_results = []
    # aggregated_results.append(f"最近{X}年，连续{M}天，top{N}的变化率")

    # 上证指数
    csv_file_path = '../data/000016.SH_daily_data.csv'
    result = calculate_top_n_changes(csv_file_path, M, N, ascending, X)
    result['index_label'] = '上证指数'
    aggregated_results.append(result)

    # 沪深300
    csv_file_path = '../data/000300.SH_daily_data.csv'
    result = calculate_top_n_changes(csv_file_path, M, N, ascending, X)
    result['index_label'] = '沪深300'
    aggregated_results.append(result)

    # 中证500
    csv_file_path = '../data/000905.SH_daily_data.csv'
    result = calculate_top_n_changes(csv_file_path, M, N, ascending, X)
    result['index_label'] = '中证500'
    aggregated_results.append(result)

    # 科创50
    csv_file_path = '../data/000688.SH_daily_data.csv'
    result = calculate_top_n_changes(csv_file_path, M, N, ascending, X)
    result['index_label'] = '科创50'
    aggregated_results.append(result)

    # 创业板
    csv_file_path = '../data/399006.SZ_daily_data.csv'
    result = calculate_top_n_changes(csv_file_path, M, N, ascending, X)
    result['index_label'] = '创业板'
    aggregated_results.append(result)

    # 合并并输出到一个文件
    combined_df = pd.concat(aggregated_results, ignore_index=True)
    combined_df.to_csv('../result/top_n_changes_all.csv', index=False, encoding='utf-8-sig')
    print("已保存所有计算结果到 ../result/top_n_changes_all.csv")
